#include "stdafx.h"
#include "PEhook.h"
#include "KeyString.h"

initialiseSingleton(CPEHook);
CPEHook::CPEHook()
{
}

CPEHook::~CPEHook()
{
    DWORD dw = 0;
    WriteProcessMemory(INVALID_HANDLE_VALUE, m_pPEAddress, &m_pHookAddress, 4, &dw);
}

void CPEHook::initData(DWORD hh)
{
    m_hDll = hh;
    PIMAGE_DOS_HEADER pp = (PIMAGE_DOS_HEADER)hh;
    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)(hh + pp->e_lfanew);
    m_hSize = nt->OptionalHeader.SizeOfImage;

    MEMORY_BASIC_INFORMATION oInfo;
    VirtualQuery(HANDLE(hh+m_hSize), &oInfo, sizeof(oInfo));
    if (0 == oInfo.AllocationBase && 0 == oInfo.AllocationProtect && oInfo.Type == 0)
    {
        m_hSize += oInfo.RegionSize;
    }
    HMODULE hNT = GetModuleHandle(KEY_STR(Key_Ntdll));
    m_pHookAddress = GetProcAddress(hNT, KEY_STRANSI(Key_ZwQueryVirtualMemory));
}

void CPEHook::hookData(HANDLE _fun)
{
    HMODULE hKer = GetModuleHandle(KEY_STR(Key_Kernaldll));
    if (hKer==0)
    {
        return;
    }

    PIMAGE_DOS_HEADER pDos = (PIMAGE_DOS_HEADER)hKer;
    PIMAGE_NT_HEADERS pNt = (PIMAGE_NT_HEADERS)((DWORD)pDos + pDos->e_lfanew);

    DWORD dwSize = pNt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].Size;
    HANDLE* pAryFun = (HANDLE*)((DWORD)hKer + pNt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress);

    for (DWORD i = 0; i < dwSize; ++i)
    {
        if (pAryFun[i]==m_pHookAddress)
        {
            m_pPEAddress = &pAryFun[i];
            DWORD dw = 0;
            DWORD dwTemp = 0;
            //VirtualProtect((LPVOID)m_pPEAddress, 4 , PAGE_EXECUTE_READWRITE,&dwTemp);
            WriteProcessMemory(INVALID_HANDLE_VALUE, m_pPEAddress, &_fun, 4, &dw);
            //VirtualProtect((LPVOID)m_pPEAddress, 4 , dwTemp,&dw);
        }
    }
}

SIZE_T CPEHook::_MyVQ( __in HANDLE hProcess, __in_opt LPCVOID lpAddress, __in DWORD dwZero, __out PMEMORY_BASIC_INFORMATION lpBuffer, __in SIZE_T dwLength, __out DWORD& pSize )
{
    typedef SIZE_T (WINAPI *_RealVQ)(HANDLE, LPCVOID, DWORD, PMEMORY_BASIC_INFORMATION, SIZE_T, DWORD&);
    _RealVQ fun = (_RealVQ)getSingleton().m_pHookAddress;
    SIZE_T Ret = fun(hProcess, lpAddress, dwZero, lpBuffer, dwLength, pSize);
    PMEMORY_BASIC_INFORMATION pBuf = (PMEMORY_BASIC_INFORMATION)lpBuffer;
    DWORD dw = (SIZE_T)pBuf->BaseAddress;
    if (dw + pBuf->RegionSize == getSingleton().m_hDll)
    {
        pBuf->RegionSize += getSingleton().m_hSize;
    }
    return Ret;
}